import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore.ops import functional as F
from mindspore.common.initializer import initializer


class TestNet(nn.Cell):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = ops.BatchNorm()

    def construct(self, input_x, scale, bias, mean, variance):
        out = self.op(input_x, scale, bias, mean, variance)
        return [out[0]]


if __name__ == "__main__":
    net = TestNet()

    input_x = ms.Tensor(np.ones([2, 2]), ms.float32)
    scale = ms.Tensor(np.ones([2]), ms.float32)
    bias = ms.Tensor(np.ones([2]), ms.float32)
    mean = ms.Tensor(np.ones([2]), ms.float32)
    variance = ms.Tensor(np.ones([2]), ms.float32)
    result = net(input_x, scale, bias, mean, variance)
    print(result[0])
    print(result[0].shape)
